# -*- coding: utf-8 -*-
"""
文件名称: summarize_across_exps.py

功能概述:
    1. 针对一个总的评估目录(eval_root), 该目录下有多个子实验子目录,
       每个子目录形如:
           sub_exp_1/
             └─ scored/
                 └─ step_metrics_summary.json
       其中 step_metrics_summary.json 为单实验的 "step -> metrics" 映射。
    2. 本脚本在所有子实验之间进行横向对比, 关注不同
       “其他题型保留比例(0%,5%,20%,50%,100%...)” 对各题型 strict_accuracy 的影响。
    3. 作图方式:
       - 每个题型/指标一张图:
         * 横轴: 训练 step
         * 纵轴: 对应 strict_accuracy
         * 一张图中多条折线, 每条线对应一个子实验(不同的数据采样比例)。
       - 背景灰色, 有网格, 折线 + 方形点, dpi=600,
         字体使用当前目录下的本地字体文件 ./SarasaMonoCL-Regular.ttf。
    4. 额外约束:
       - 全局变量 MAX_STEP = 1000:
         * 只考虑 step <= MAX_STEP 的数据点。
         * 超过 MAX_STEP 的 step 在汇总和作图中一律忽略。

输出:
    - 一个整理后的总 JSON:
      out_dir/summary_across_exps_all_steps.json
    - 若干图像:
      out_dir/overall_strict_accuracy.png
      out_dir/judge_strict_accuracy.png
      out_dir/compare_strict_accuracy.png
      out_dir/diagnosis_strict_accuracy.png
      out_dir/choice_strict_accuracy.png

使用示例:
    python fundus_reasoner/eval/summarize_across_exps.py \
        --eval-root experiments/eval_output \
        --out-dir experiments/eval_output/summary_across_exps
"""

import argparse
import json
import logging
from pathlib import Path
from typing import Dict, Any, List, Tuple

import matplotlib
import matplotlib.pyplot as plt

# 全局字体路径: 默认使用当前工作目录下的本地字体文件
FONT_PATH = Path("./SarasaMonoCL-Regular.ttf")

# 全局 step 上限: 只考虑 step <= MAX_STEP 的数据
MAX_STEP = 1000


logger = logging.getLogger(__name__)
if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter(
        fmt="[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)


# -----------------------------
# 字体设置: 使用本地 ./SarasaMonoCL-Regular.ttf
# -----------------------------
def setup_font() -> None:
    """使用当前目录下的本地字体文件 ./SarasaMonoCL-Regular.ttf

    要求:
        - FONT_PATH 指向实际存在的 .ttf 文件
        - 若文件不存在或加载失败则抛异常, 防止悄悄退回系统默认字体导致中文渲染异常
    """
    font_file = FONT_PATH

    if not font_file.is_file():
        raise FileNotFoundError(
            f"字体文件不存在: {font_file} "
            f"（请确保 SarasaMonoCL-Regular.ttf 放在当前工作目录, "
            f"或修改 summarize_across_exps.py 中的 FONT_PATH）"
        )

    try:
        import matplotlib.font_manager as fm

        # 显式注册该字体文件
        fm.fontManager.addfont(str(font_file))
        # 通过文件获取字体名称, 避免硬编码 family 名字
        font_name = fm.FontProperties(fname=str(font_file)).get_name()

        matplotlib.rcParams["font.family"] = font_name
        matplotlib.rcParams["axes.unicode_minus"] = False

        logger.info("成功加载本地字体: %s, 字体名: %s", font_file, font_name)
    except Exception as e:
        raise RuntimeError(
            f"加载字体文件失败: {font_file}, error={e}"
        ) from e


# -----------------------------
# 工具函数: 读取每个子实验的 summary JSON
# -----------------------------
def load_experiment_summary(exp_dir: Path) -> Dict[str, Any]:
    """从子实验目录中读取 scored/step_metrics_summary.json

    要求:
        - exp_dir 是一个包含 scored 子目录的目录
        - scored/step_metrics_summary.json 必须存在, 否则报错
    返回:
        data: dict, 形如 { "1304": { "overall_strict_accuracy": ..., ... }, ... }
    """
    summary_path = exp_dir / "scored" / "step_metrics_summary.json"
    if not summary_path.is_file():
        raise RuntimeError(
            f"子实验目录缺少 scored/step_metrics_summary.json: {summary_path}"
        )

    try:
        data = json.loads(summary_path.read_text(encoding="utf-8"))
    except Exception as e:
        raise RuntimeError(
            f"读取或解析子实验 summary JSON 失败: path={summary_path}, error={e}"
        ) from e

    if not isinstance(data, dict):
        raise RuntimeError(
            f"summary JSON 顶层结构必须为 dict (step->metrics), 实际为: {type(data)}"
        )

    return data


# -----------------------------
# 工具函数: 解析子目录名 -> ratio, display_name
# -----------------------------
def parse_exp_name(name: str) -> Tuple[int, str]:
    """根据子目录名解析 “其他题型保留比例” ratio 及中文显示名

    约定:
        - 以 "_diag_only" 结尾 -> 仅诊断题, ratio=0
        - 以 "_pXX" 结尾      -> 其他题型保留 XX%, ratio=XX
        - 其他情况            -> 认为其他题型全部保留, ratio=100

    返回:
        (ratio, display_name)

    display_name 约定:
        - ratio=0  -> "仅诊断题 (其他题型 0%)"
        - ratio>0  -> "诊断题 + 其他题型 {ratio}%"
    """
    ratio: int = 100

    if name.endswith("_diag_only"):
        ratio = 0
    else:
        parts = name.split("_")
        last = parts[-1]
        if last.startswith("p") and last[1:].isdigit():
            ratio = int(last[1:])
        else:
            ratio = 100

    if ratio == 0:
        display = "仅诊断题 (其他题型 0%)"
    else:
        display = f"诊断题 + 其他题型 {ratio}%"

    return ratio, display


# -----------------------------
# 工具函数: 汇总所有子实验为一个大结构
# -----------------------------
def gather_all_experiments(
    eval_root: Path,
) -> Dict[str, Dict[str, Any]]:
    """在 eval_root 下汇总所有子实验的 summary

    只保留 step <= MAX_STEP 的数据点。

    返回:
        experiments = {
          exp_name: {
            "ratio": int,
            "display_name": str,
            "steps": {
                step(int): metrics(dict)
            }
          },
          ...
        }

    若 eval_root 下任一子目录缺少 summary 文件, 直接报错。
    若某子实验在过滤 step > MAX_STEP 后不剩任何 step, 也报错。
    """
    if not eval_root.is_dir():
        raise RuntimeError(f"评估根目录不存在或不是目录: {eval_root}")

    subdirs = [d for d in eval_root.iterdir() if d.is_dir()]
    if not subdirs:
        raise RuntimeError(f"评估根目录下未发现任何子实验子目录: {eval_root}")

    logger.info("在评估根目录 %s 下发现 %d 个子实验目录", eval_root, len(subdirs))
    logger.info("全局 step 上限 MAX_STEP=%d, 将忽略所有 step > MAX_STEP 的数据", MAX_STEP)

    experiments: Dict[str, Dict[str, Any]] = {}

    for d in subdirs:
        exp_name = d.name
        ratio, display_name = parse_exp_name(exp_name)
        data = load_experiment_summary(d)  # {step_str: metrics}

        steps_map: Dict[int, Dict[str, Any]] = {}
        for k, v in data.items():
            try:
                step_int = int(k)
            except Exception:
                raise RuntimeError(
                    f"子实验 {exp_name} 的 summary 中存在无法转换为 int 的 step 字符串: {k!r}"
                )
            if step_int > MAX_STEP:
                # 超过全局上限的 step 直接跳过
                continue
            if not isinstance(v, dict):
                raise RuntimeError(
                    f"子实验 {exp_name} 的 step={k} 对应条目不是 dict 类型, 实际: {type(v)}"
                )
            steps_map[step_int] = v

        if not steps_map:
            raise RuntimeError(
                f"子实验 {exp_name} 在 step <= {MAX_STEP} 的范围内没有任何数据点, "
                f"请检查是否 step 全部在 {MAX_STEP} 之后。"
            )

        experiments[exp_name] = {
            "ratio": ratio,
            "display_name": display_name,
            "steps": steps_map,
        }

    return experiments


# -----------------------------
# 工具函数: 画图 (每个指标一张, 多条折线)
# -----------------------------
def plot_metrics_across_experiments(
    experiments: Dict[str, Dict[str, Any]],
    out_dir: Path,
) -> None:
    """针对所有 step, 以 step 为横轴, 以 strict_accuracy 为纵轴, 按题型作图

    参数:
        experiments: gather_all_experiments 的返回值
        out_dir:     图像输出目录, 自动创建

    说明:
        - 每个指标一张图:
          * X 轴: 训练 step (整数)
          * Y 轴: 指标值 (strict_accuracy)
          * 一张图中多条折线, 每条线对应一个子实验(不同采样比例)。
        - 只考虑 step <= MAX_STEP 的数据 (在 gather_all_experiments 中已过滤)。
        - 如果某个实验在某些 step 上没有指标数据, 则该实验的折线在这些 step 上没有点。
    """
    out_dir.mkdir(parents=True, exist_ok=True)

    # 收集所有出现过的 step, 用于决定 x 轴范围
    all_steps: List[int] = []
    for info in experiments.values():
        all_steps.extend(info["steps"].keys())
    all_steps = sorted(set(all_steps))

    if not all_steps:
        raise RuntimeError(
            f"在所有子实验中 (且 step <= {MAX_STEP}) 未找到任何 step, 无法作图"
        )

    metric_keys = [
        "overall_strict_accuracy",
        "judge_strict_accuracy",
        "compare_strict_accuracy",
        "diagnosis_strict_accuracy",
        "choice_strict_accuracy",
    ]

    # 给每个实验分配颜色: 按 ratio 排序, 用 matplotlib 默认色循环
    exp_names_sorted = sorted(
        experiments.keys(),
        key=lambda name: experiments[name]["ratio"],
    )
    color_cycle = plt.rcParams["axes.prop_cycle"].by_key().get("color", [])
    if not color_cycle:
        color_cycle = ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b"]

    exp_colors: Dict[str, str] = {}
    for idx, name in enumerate(exp_names_sorted):
        exp_colors[name] = color_cycle[idx % len(color_cycle)]

    for mk in metric_keys:
        plt.figure(figsize=(10, 6))
        ax = plt.gca()
        ax.set_facecolor("#f0f0f0")
        ax.grid(True, linestyle="--", alpha=0.5)

        # 每个实验一条线
        for exp_name in exp_names_sorted:
            info = experiments[exp_name]
            ratio = info["ratio"]
            display_name = info["display_name"]
            steps_map = info["steps"]  # Dict[int, Dict[str, Any]]

            steps_sorted = sorted(steps_map.keys())
            x_vals: List[int] = []
            y_vals: List[float] = []
            for s in steps_sorted:
                metrics = steps_map[s]
                if mk not in metrics:
                    continue
                try:
                    y = float(metrics[mk])
                except Exception:
                    continue
                x_vals.append(s)
                y_vals.append(y)

            if not x_vals:
                logger.warning(
                    "实验 %s 在指标 %s 上 (且 step <= %d) 无有效数据, 不绘制该折线",
                    exp_name,
                    mk,
                    MAX_STEP,
                )
                continue

            color = exp_colors[exp_name]
            label = f"{display_name} (ratio={ratio}%)"

            ax.plot(
                x_vals,
                y_vals,
                marker="s",
                linestyle="-",
                linewidth=2,
                markersize=3,
                color=color,
                label=label,
            )

        ax.set_xlabel("训练 step", fontsize=12)
        ax.set_ylabel("strict_accuracy", fontsize=12)
        ax.set_title(
            f"{mk} 随训练步数变化曲线 (只展示 step ≤ {MAX_STEP})",
            fontsize=14,
        )

        # x 轴范围: 在有效 step 范围内, 但不超过 MAX_STEP
        xmin = min(all_steps)
        xmax = min(max(all_steps), MAX_STEP)
        ax.set_xlim(xmin, xmax)

        ax.legend(loc="best", fontsize=9)

        plt.tight_layout()

        out_png = out_dir / f"{mk}.png"
        try:
            plt.savefig(out_png, dpi=600)
            logger.info("指标图已保存: %s", out_png)
        except Exception as e:
            raise RuntimeError(f"保存指标图失败: path={out_png}, error={e}") from e
        finally:
            plt.close()


# -----------------------------
# 主逻辑: 汇总 + 画图
# -----------------------------
def summarize_across_experiments(
    eval_root: Path,
    out_dir: Path,
) -> None:
    """在 eval_root 下汇总多个子实验, 并按题型作图 (横轴 step, 纵轴分值, 多条折线)

    只考虑 step <= MAX_STEP 的数据点。
    """

    setup_font()

    experiments = gather_all_experiments(eval_root)

    # 输出整体汇总 JSON, 方便后续查看或二次分析
    out_dir.mkdir(parents=True, exist_ok=True)
    summary_obj = {
        "max_step": MAX_STEP,
        "experiments": {
            exp_name: {
                "ratio": info["ratio"],
                "display_name": info["display_name"],
                "steps": {
                    str(step): metrics
                    for step, metrics in sorted(info["steps"].items(), key=lambda x: x[0])
                },
            }
            for exp_name, info in experiments.items()
        },
    }
    summary_path = out_dir / "summary_across_exps_all_steps.json"
    try:
        summary_path.write_text(
            json.dumps(summary_obj, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("跨实验汇总 JSON 已写入: %s", summary_path)
    except Exception as e:
        raise RuntimeError(
            f"写入 summary_across_exps_all_steps.json 失败: path={summary_path}, error={e}"
        ) from e

    # 绘图: 每个指标一张图, 一张图里多条折线
    plot_metrics_across_experiments(
        experiments=experiments,
        out_dir=out_dir,
    )


# -----------------------------
# CLI
# -----------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="跨子实验汇总 step_metrics_summary.json, 以 step 为横轴、分值为纵轴按题型作图",
    )
    parser.add_argument(
        "--eval-root",
        "-i",
        type=str,
        required=True,
        help="评估根目录, 其下每个子目录为一个子实验, 且含 scored/step_metrics_summary.json",
    )
    parser.add_argument(
        "--out-dir",
        "-o",
        type=str,
        default=None,
        help="汇总 JSON 与图像输出目录 (默认: <eval-root>/summary_across_exps)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    eval_root = Path(args.eval_root)
    if args.out_dir is not None:
        out_dir = Path(args.out_dir)
    else:
        out_dir = eval_root / "summary_across_exps"

    summarize_across_experiments(
        eval_root=eval_root,
        out_dir=out_dir,
    )


if __name__ == "__main__":
    main()
